
#include "reduce_log_sum_exp_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <algorithm>
#include <vector>

using namespace std;
namespace optiling {

    const int32_t BLOCK_SIZE = 32;
    struct TilingArg {
        int32_t tileNum;
        int32_t tileLength;
        int32_t lastTileLength;
    };
    // 获得外层组数
    int32_t getGroupNum(vector<int32_t> &x_shape, int32_t dim){
        int32_t groupNum = 1;
        for(int32_t i = 0; i < dim; ++i){
            groupNum *= x_shape[i];
        }
        return groupNum;
    }

    // 获得数据总个数
    int32_t getTotalLength(vector<int32_t> &x_shape){
        int32_t totalLength = 1;
        for(int32_t i = 0; i < x_shape.size(); ++i){
            totalLength *= x_shape[i];
        }
        return totalLength;
    }

    int32_t getTotalLengthAligned(int32_t totalLength, int32_t sizeOfDatatype){
        int32_t totalLengthAligned;
        uint32_t ALIGN_NUM = BLOCK_SIZE / sizeOfDatatype;
        if (totalLength % ALIGN_NUM != 0) {  //不对齐，先32位对齐
            totalLengthAligned = ((totalLength + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
        } else {
            totalLengthAligned = totalLength;
        }
        return totalLengthAligned;
    }

    // 获得组长度
    int32_t getGroupLength(int32_t totalLength, int32_t groupNum){
        return totalLength / groupNum;
    }

    // 获得消除的向量个数
    int32_t getXiNum(vector<int32_t> &x_shape, int32_t dim){
        return x_shape[dim];
    }

    // 获得 xi 的长度
    int32_t getXiLength(int32_t groupLength, int32_t xiNum){
        return groupLength / xiNum;
    }

    // 获得对齐的 xiLength
    int32_t getXiLengthAligned(int32_t xiLength, int32_t sizeOfDatatype){
        int32_t xiLengthAligned;
        uint32_t ALIGN_NUM = BLOCK_SIZE / sizeOfDatatype;
        if (xiLength % ALIGN_NUM != 0) {  //不对齐，先32位对齐
            xiLengthAligned = ((xiLength + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
        } else {
            xiLengthAligned = xiLength;
        }
        return xiLengthAligned;
    }

    // 获得 tile 参数，单 buffer
    TilingArg getTilingArg(int32_t lengthAligned, uint64_t ub_block_num, int32_t sizeOfDatatype){
        int32_t tileNum;
        int32_t tileLength = 0;
        int32_t lastTileLength = 0;
        uint32_t ALIGN_NUM = BLOCK_SIZE / sizeOfDatatype;  //每一块里面的数据个数

        tileNum = lengthAligned / ALIGN_NUM / ub_block_num;  //总长度/(ub的块个数*每一块里面的数据个数)
        if (tileNum == 0) {
            tileNum = 1;
            tileLength = lengthAligned;
            lastTileLength = tileLength;
        } else if((lengthAligned / ALIGN_NUM) % ub_block_num == 0){
            tileLength = ub_block_num * ALIGN_NUM;
            lastTileLength = tileLength;
        }else{
            tileNum = tileNum + 1;
            tileLength = ub_block_num * ALIGN_NUM;
            lastTileLength = lengthAligned - (tileNum - 1) * tileLength;
        }
        TilingArg arg;
        arg.tileNum = tileNum;
        arg.tileLength = tileLength;
        arg.lastTileLength = lastTileLength;

        return arg;
    }

    void printShape(vector<int32_t> &x_shape){
        printf("shape: ");
        for(int i = 0; i < x_shape.size(); ++i){
            printf("%d ", x_shape[i]);
        }
        printf("\n");
    }


// static ge::graphStatus TilingFunc(gert::TilingContext* context)
// {

//   ReduceLogSumExpTilingData tiling;

//   const gert::StorageShape* x1_shape = context->GetInputShape(0);


//   auto shape_x = context->GetInputTensor(0)->GetOriginShape();
//   int32_t x_ndarray[4];
//   int32_t x_dimensional;
//   int32_t size = 1;

//   x_dimensional = shape_x.GetDimNum();
//   for(int i = 0; i < x_dimensional; i++)
//   {
//       x_ndarray[i] = shape_x.GetDim(i);
//       size *= x_ndarray[i];
//   }
//   tiling.set_size(size);
//   tiling.set_x_ndarray(x_ndarray);
//   tiling.set_x_dimensional(x_dimensional);

//   bool keep = *context->GetAttrs()->GetBool(0);
//   tiling.set_keep(keep);

//   uint64_t ubSize;
//   auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
//   ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); //获取硬件平台存储空间 UB 的内存大小
//   int32_t tileDataMaxNum = ubSize / 32 / 4 / 5 * 32;// *6;
//   tiling.set_tileDataMaxNum(tileDataMaxNum);//(5*1024);//(30*1024); //每次处理的数据量


static ge::graphStatus TilingFunc(gert::TilingContext* context) {
    ReduceLogSumExpTilingData tiling;

    // 1. 获取平台信息
    uint64_t ub_size;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto socVersion = ascendcPlatform.GetSocVersion();
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
    auto aivNum = 1;  // vector core num  1
    // 2. 获取数据信息
    auto dt = context->GetInputTensor(0)->GetDataType();
    // 获取 shape 信息
    gert::Shape originShape = context->GetInputTensor(0)->GetOriginShape();
    int32_t  shape_size = originShape.GetDimNum();
    vector<int32_t> x_shape(shape_size);
    for(int i = 0; i < shape_size; ++i){
        x_shape[i] = originShape[i];
    }
    printShape(x_shape);
    // 得到轴信息
    auto addr = context->GetInputTensor(1)->GetData<int64_t>();
    int32_t axes_shape = addr[0];
    //printf("axes_shape: %d ", axes_shape); //要删除的维度位置                     
    if(axes_shape < 0){                    //处理-1情况后的维度的位置
        axes_shape = 0;
    }else{
        axes_shape = axes_shape;
    }
    //printf("axes_shape: %d ", axes_shape); //要删除的维度位置
    //printf("\n");

    if(dt == ge::DT_FLOAT){
        //printf("float\n");
        context->SetTilingKey(1);
        // 计算每个 dim 对应的参数
        uint32_t data_num = 2;
        uint32_t ub_block_num = ub_size / BLOCK_SIZE / data_num;
        TilingArg args;
        int32_t sizeOfDatatype = 4;
        int32_t groupNum, groupLength;
        int32_t xiNum, xiLength, xiLengthAligned;
        int32_t tileNum, tileLength, lastTileLength;
        int32_t totalLength, totalLengthAligned;
        int32_t cnt = 0;

        groupNum = getGroupNum(x_shape, axes_shape);
        totalLength = getTotalLength(x_shape);
        totalLengthAligned = getTotalLengthAligned(totalLength, sizeOfDatatype);
        groupLength = getGroupLength(totalLength, groupNum);
        xiNum = getXiNum(x_shape, axes_shape);
        xiLength = getXiLength(groupLength, xiNum);
        xiLengthAligned = getXiLengthAligned(xiLength, sizeOfDatatype);
        args =  getTilingArg(xiLengthAligned, ub_block_num, sizeOfDatatype);
        tileNum = args.tileNum;
        tileLength = args.tileLength;
        lastTileLength = args.lastTileLength;

        // printf("------------------------------------------\n");
        // printShape(x_shape);
        // printf("coreLength = %d\n", totalLengthAligned);
        // printf("groupNum = %d\n", groupNum);
        // printf("groupLength = %d\n", groupLength);
        // printf("xiNum = %d\n", xiNum);
        // printf("xiLength = %d\n", xiLength);
        // printf("xiLengthAligned = %d\n", xiLengthAligned);
        // printf("tileNum = %d\n", tileNum);
        // printf("tileLength = %d\n", tileLength);
        // printf("lastTileLength = %d\n", lastTileLength);
        
        tiling.set_coreLength(totalLengthAligned);
        tiling.set_groupNum(groupNum);
        tiling.set_groupLength(groupLength);
        tiling.set_xiNum(xiNum);
        tiling.set_xiLength(xiLength);
        tiling.set_xiLengthAligned(xiLengthAligned);
        tiling.set_tileNum(tileNum);
        tiling.set_tileLength(tileLength);
        tiling.set_lastTileLength(lastTileLength);
    }else{
        //printf("half\n");
        context->SetTilingKey(2);
        // 计算每个 dim 对应的参数
        uint32_t data_num = 7;
        uint32_t ub_block_num = ub_size / BLOCK_SIZE / data_num;
        TilingArg args;
        int32_t sizeOfDatatype = 2;
        int32_t groupNum, groupLength;
        int32_t xiNum, xiLength, xiLengthAligned;
        int32_t tileNum, tileLength, lastTileLength;
        int32_t totalLength, totalLengthAligned;

        groupNum = getGroupNum(x_shape, axes_shape);
        totalLength = getTotalLength(x_shape);
        totalLengthAligned = getTotalLengthAligned(totalLength, sizeOfDatatype);
        groupLength = getGroupLength(totalLength, groupNum);
        xiNum = getXiNum(x_shape, axes_shape);
        xiLength = getXiLength(groupLength, xiNum);
        xiLengthAligned = getXiLengthAligned(xiLength, sizeOfDatatype);
        args =  getTilingArg(xiLengthAligned, ub_block_num, sizeOfDatatype);
        tileNum = args.tileNum;
        tileLength = args.tileLength;
        lastTileLength = args.lastTileLength;

        // printf("------------------------------------------\n");
        // printShape(x_shape);
        // printf("coreLength = %d\n", totalLengthAligned);
        // printf("groupNum = %d\n", groupNum);
        // printf("groupLength = %d\n", groupLength);
        // printf("xiNum = %d\n", xiNum);
        // printf("xiLength = %d\n", xiLength);
        // printf("xiLengthAligned = %d\n", xiLengthAligned);
        // printf("tileNum = %d\n", tileNum);
        // printf("tileLength = %d\n", tileLength);
        // printf("lastTileLength = %d\n", lastTileLength);
        
        tiling.set_coreLength(totalLengthAligned);
        tiling.set_groupNum(groupNum);
        tiling.set_groupLength(groupLength);
        tiling.set_xiNum(xiNum);
        tiling.set_xiLength(xiLength);
        tiling.set_xiLengthAligned(xiLengthAligned);
        tiling.set_tileNum(tileNum);
        tiling.set_tileLength(tileLength);
        tiling.set_lastTileLength(lastTileLength);
    }

  context->SetBlockDim(1);
  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

  return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x_shape;
    return GRAPH_SUCCESS;
}
static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
{
const auto inputDataType = context->GetInputDataType(0);
context->SetOutputDataType(0, inputDataType);
return ge::GRAPH_SUCCESS;
}
}


namespace ops {
class ReduceLogSumExp : public OpDef {
public:
    explicit ReduceLogSumExp(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("axes")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("keep_dims").AttrType(OPTIONAL).Bool(false);

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(ReduceLogSumExp);
}
